# -*- coding: utf-8 -*-
"""
@file name      : hook_fmap_vis.py
@author         : QuZhang
@date           : 2020-12-30 22:40
@brief          : 使用hook可视化特征图
"""
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import transforms
from PIL import Image
import torchvision.models as models
import torch.nn as nn
import numpy as np
import torchvision.utils as vutils
import torch


# ----------- feature map visualization --------------
flag = True
if flag:
    # writer = SummaryWriter(comment='vis_alexnet', filename_suffix="alexnet_fmap")
    writer = SummaryWriter(comment="vis_alexnet_inplace", filename_suffix='inplace_Fasle')

    # 数据
    path_img = "./lena.png"
    normMean = [0.49139968, 0.48215827, 0.44653124]
    normStd = [0.49139968, 0.48215827, 0.44653124]

    img_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=normMean, std=normStd),
    ])

    img_pil = Image.open(path_img).convert("RGB")
    if img_transform is not None:
        img_tensor = img_transform(img_pil)
    img_tensor.unsqueeze_(0)  # chw --> bchw

    # 模型
    alexnet = models.alexnet(pretrained=True)

    # 注册hook
    fmap_dict = dict()
    # 遍历模型的每一个子模块/每一层
    for name, sub_module in alexnet.named_modules():

        if isinstance(sub_module, nn.Conv2d):
            key_name = str(sub_module.weight.shape)  # 获取卷积层形状
            fmap_dict.setdefault(key_name, list())  # 将卷积层的形状作为键, 实值默认为列表

            n1, n2 = name.split('.')  # 用.划分字符串,比如linear1.conv1变为linear1和conv1

            def hook_func(m, i, o):
                """前向传播时，自动调用回调函数
                """
                key_name = str(m.weight.shape)
                fmap_dict[key_name].append(o)

            # 获取卷积层
            alexnet._modules[n1]._modules[n2].register_forward_hook(hook_func)  # 为卷积层注册回调函数hook_func

        # forward
        output = alexnet(img_tensor)

        # add image
        with torch.no_grad():
            for layer_name, fmap_list in fmap_dict.items():
                print("layer name: ", layer_name)
                print("fmap_list size: ", len(fmap_list))
                fmap = fmap_list[0]
                print("fmap[0] shape: {}".format(fmap.shape))
                fmap.transpose_(0, 1) # bchw -> cbhw

                nrow = int(np.sqrt(fmap.shape[0]))
                fmap_grid = vutils.make_grid(fmap, normalize=True, scale_each=True, nrow=nrow)
                writer.add_image('feature map in {}'.format(layer_name), fmap_grid, global_step=322)
